import base64
import os
import uuid
import json
import time
import jwt
import requests
import cv2
from typing import Optional, Tuple, Dict, Any

import urllib3
from dotenv import load_dotenv

from utils import sr_s3
from utils.sr_utils import image_path_to_data_url


class videoGenerator_kling:
    """视频生成工具类，封装KlingAI视频生成API调用逻辑，支持外部调用"""

    def __init__(self):
        """初始化配置与环境变量"""
        # 加载环境变量
        self.verify_ssl = True
        load_dotenv(verbose=True, override=True)
        self.app_env = os.getenv('APP_ENV', 'prod')
        self.KELING_URL = os.getenv('KELING_URL')
        self.KELING_AK = os.getenv('KELING_AK')
        self.KELING_SK = os.getenv('KELING_SK')
        self.MODELS = {
            1: "kling-v1",
            2: "kling-v1-5",
            3: "kling-v1-6",
            4: "kling-v2-master",
            5: "kling-v2-1",
            6: "kling-v2-1-master",
        }
        # 初始化鉴权头（JWT动态生成，30分钟有效期）
        self.headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self._encode_jwt_token()}"
        }

        # 临时文件存储目录（可外部配置）
        self.temp_image_dir = "./upload"
        self.temp_video_dir = "./download"
        os.makedirs(self.temp_image_dir, exist_ok=True)
        os.makedirs(self.temp_video_dir, exist_ok=True)

    # --------------------------
    # 内部工具方法：鉴权与配置
    # --------------------------
    def _encode_jwt_token(self) -> str:
        """生成JWT鉴权令牌（有效期30分钟）"""
        headers = {"alg": "HS256", "typ": "JWT"}
        payload = {
            "iss": self.KELING_AK,
            "exp": int(time.time()) + 1800,  # 30分钟后过期
            "nbf": int(time.time()) - 5  # 5秒前生效（避免时间偏差）
        }
        tokens = jwt.encode(payload, self.KELING_SK, headers=headers)
        return tokens

    # def _encode_jwt_token(self):
    #     headers = {
    #         "alg": "HS256",
    #         "typ": "JWT"
    #     }
    #     payload = {
    #         "iss": "Ad8QPQKKaF8FYHBCp8kGRHAbKDtm4BAp",
    #         "exp": int(time.time()) + 1800,  # 有效时间，此处示例代表当前时间+1800s(30min)
    #         "nbf": int(time.time()) - 5  # 开始生效的时间，此处示例代表当前时间-5秒
    #     }
    #     token = jwt.encode(payload, "Ad8QPQKKaF8FYHBCp8kGRHAbKDtm4BAp", headers=headers)
    #     return token

    def refresh_auth_token(self) -> None:
        """刷新鉴权令牌（适用于长时任务）"""
        self.headers["Authorization"] = f"Bearer {self._encode_jwt_token()}"

    # --------------------------
    # 内部工具方法：文件处理
    # --------------------------
    def _delete_temp_files(self, *file_paths: str) -> None:
        """删除临时文件（忽略不存在的文件）"""
        for path in file_paths:
            if path and os.path.exists(path):
                try:
                    os.remove(path)
                except Exception as e:
                    print(f"清理临时文件失败 {path}: {str(e)}")

    def _download_and_compress_image(self, img_url: str, storage_service: str, bucket_name: str) -> Tuple[bool, str, bool]:
        """下载图片并压缩（返回是否成功、base64编码、是否压缩）"""
        try:
            # 生成临时保存路径
            name = str(uuid.uuid4())
            ext = os.path.splitext(img_url)[-1].lower() or '.jpg'
            save_path = os.path.join(self.temp_image_dir, f"{name}{ext}")

            # 调用S3工具下载图片（假设sr_s3为外部存储工具）
            from utils import sr_s3  # 延迟导入，避免外部依赖问题
            success, err = sr_s3.download_file(
                img_url, save_path, serviceName=storage_service, bucketName=bucket_name
            )
            if not success:
                return False, f"下载失败: {err}", False
            # save_path = './image/content_2869596861640069.jpg'
            # 压缩图片（确保符合模型要求）
            from utils.sr_utils import video_compress_image_size_path
            success, new_path = video_compress_image_size_path(save_path, 10, (300, 6000), (300, 6000))
            if not success:
                return False, new_path, False
                # 读取临时文件并转为base64
            with open(new_path, 'rb') as f:
                base64_data = base64.b64encode(f.read()).decode('utf-8')
            # os.remove(new_path)  # 清理临时文件
            return True, base64_data, True
        except Exception as e:
            return False, f"图片处理异常: {str(e)}", False
        # finally:
        #     self._delete_temp_files(save_path)

    def _download_video(self, video_url: str, task_id: str) -> Tuple[bool, Dict[str, Any], str]:
        """下载生成的视频到本地（返回是否成功、视频信息）"""
        try:
            filename = f"video_{task_id}.mp4"
            save_path = os.path.join(self.temp_video_dir, filename)

            # 流式下载视频
            response = requests.get(video_url, stream=True, timeout=30)
            response.raise_for_status()

            with open(save_path, "wb") as f:
                for chunk in response.iter_content(chunk_size=1024 * 1024):  # 1MB/块
                    if chunk:
                        f.write(chunk)

            # 获取视频信息（宽高、帧率等）
            cap = cv2.VideoCapture(save_path)
            video_info = {
                "save_path": save_path,
                "width": int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)),
                "height": int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)),
                "fps": cap.get(cv2.CAP_PROP_FPS),
                "duration": cap.get(cv2.CAP_PROP_FRAME_COUNT) / cap.get(cv2.CAP_PROP_FPS) if cap.get(cv2.CAP_PROP_FPS) > 0 else 0
            }
            cap.release()
            return True, video_info, save_path
        except Exception as e:
            return False, {"error": f"视频下载失败: {str(e)}"}, ''

    def _upload_video_to_s3(self, video_path: str, s3_storage: str, storage_service: str, bucket_name: str) -> Tuple[bool, str]:
        """上传视频到S3（返回是否成功、S3路径）"""
        if not video_path or not os.path.exists(video_path):
            return False, "视频文件不存在"
        try:
            from utils import sr_s3
            s3_key = sr_s3.upload_login_or_visitor(
                video_path, s3_storage, serviceName=storage_service, bucketName=bucket_name
            )
            return True, s3_key
        except Exception as e:
            return False, f"上传S3失败: {str(e)}"

    # --------------------------
    # 核心方法：API调用与任务处理
    # --------------------------
    def _create_keling_video_task(self, payload: dict, api_path: str) -> Tuple[bool, str, str]:
        """调用KlingAI API创建视频任务（返回是否成功、任务ID、错误信息）"""
        try:
            url = f"{self.KELING_URL}{api_path}"
            response = requests.post(url, headers=self.headers, json=payload)
            response_data = response.json()
            print(response_data)
            if response_data.get("code") == 0:
                task_id = response_data.get("data", {}).get("task_id")
                print("task_id", task_id)
                if task_id:
                    return True, task_id, ""
                else:
                    return False, "", "API返回成功但无task_id"
            else:
                return False, "", f"API错误: {response_data.get('message', '未知错误')}（code: {response_data.get('code')}）"
        except Exception as e:
            return False, "", f"创建任务失败: {str(e)}"

    def _poll_task_status(self, task_id: str, api_path: str, timeout: int = 600, interval: int = 5) -> Tuple[bool, str]:
        """轮询任务状态（返回是否成功、视频URL/错误信息）"""
        if not task_id:
            return False, "任务ID为空"

        url = f"{self.KELING_URL}{api_path}/{task_id}"
        start_time = time.time()
        retries = 0
        max_retries = 5  # 最大重试次数
        # 创建带重试机制的Session
        session = requests.Session()
        adapter = requests.adapters.HTTPAdapter(
            max_retries=max_retries,
            pool_connections=10,
            pool_maxsize=10
        )
        session.mount('https://', adapter)
        # 禁用不安全请求的警告
        if not self.verify_ssl:
            urllib3.disable_warnings()

        while time.time() - start_time < timeout:
            try:
                # 发送请求（增加超时设置和SSL验证控制）
                response = session.get(
                    url,
                    headers=self.headers,
                    timeout=(5, 15),  # (连接超时, 读取超时)
                    verify=self.verify_ssl  # 控制是否验证SSL证书
                )
                response.raise_for_status()
                response_data = response.json()
                print(f"轮询响应: {response_data}")
                code = response_data.get("code")
                task_status = response_data.get("data", {}).get("task_status")
                message = response_data.get("data", {}).get("task_status_msg")

                if code != 0:
                    return False, f"查询失败: {message}（code: {code}）"

                if task_status == "succeed":
                    videos = response_data.get("data", {}).get("task_result", {}).get("videos", [])
                    if videos and "url" in videos[0]:
                        return True, videos[0]["url"]
                    else:
                        return False, "任务成功但无视频URL"
                elif task_status == "failed":
                    return False, message
                else:
                    # 任务处理中，继续轮询
                    print(f"任务状态: {task_status}，等待{interval}秒后重试")
                    time.sleep(interval)
                    retries = 0  # 重置重试计数

            except requests.exceptions.SSLError as e:
                retries += 1
                if retries >= max_retries:
                    return False, f"SSL错误（重试{retries}次后）: {str(e)}"
                print(f"SSL连接错误（{retries}/{max_retries}）: {str(e)}，{interval}秒后重试...")
                time.sleep(interval * retries)  # 指数退避策略

            except requests.exceptions.RequestException as e:
                retries += 1
                if retries >= max_retries:
                    return False, f"网络请求错误（重试{retries}次后）: {str(e)}"
                print(f"网络请求错误（{retries}/{max_retries}）: {str(e)}，{interval}秒后重试...")
                time.sleep(interval * retries)  # 指数退避策略

            except json.JSONDecodeError:
                retries += 1
                if retries >= max_retries:
                    return False, f"响应格式错误（重试{retries}次后）: 无法解析JSON"
                print(f"响应格式错误（{retries}/{max_retries}）: 无法解析JSON，{interval}秒后重试...")
                time.sleep(interval * retries)

            except Exception as e:
                return False, f"轮询异常: {str(e)}"

        return False, f"任务超时（超过{timeout}秒）"

    # 成功的队列写死的
    def _send_result_to_queue_success(self, queue_name, body, status, msg, results=None, start_time=None, local_time=None, local_video=None, ext="mp4"):
        msg_body = {
            'QueueResult': queue_name,
            'status': status,
            'taskId': body.get('taskId'),
            'taskType': body.get('taskType'),
            'appId': body.get('appId'),
            'machineId': 1,
            'width': 960,
            'height': 960,
            'tokens': 108900,
            'fps': 24,
            'results': {'key': 'video/visitor/image_to_video/2025-07-17/549152e6660d496187a98b1bcdc70fe9/video_task_cgt-20250717101955-86c56.mp4'},
            'msgFetchEndTime': start_time,
            'start_time': start_time,
            'ext': ext,
            'msg': msg
        }
        return msg_body

    # --------------------------
    # 外部调用方法：主入口
    # --------------------------
    def process_video_task(self, task_params: Dict[str, Any]):
        """
        处理视频生成任务（外部调用入口）

        参数:
            task_params: 任务参数字典，包含：
                - taskId: 任务ID
                - modelName: 模型名称索引
                - prompt: 提示词
                - modelType: 模型类型（1:文生视频, 2:单图生视频, 3:首尾图生视频, 4:单图特效, 5:多图特效）
                - duration: 视频时长（默认5）
                - first_image: 首图URL（modelType=2/3/4/5时需要）
                - last_image: 尾图URL（modelType=3/5时需要）
                - storageService: 存储服务类型（如'r2'）
                - bucketName: 存储桶名称
                - storedPrefix: S3存储路径前缀
                - QueueResult: 结果队列名称
                - 其他可选参数：fps, resolution, ratio等

        返回:
            (是否成功, 结果字典) 结果字典包含's3_key'（视频S3路径）或'error'（错误信息）
        """
        # 初始化变量
        start_time = time.time()
        task_id = task_params.get('taskId', f"temp_{uuid.uuid4().hex[:8]}")
        first_image_path = None
        last_image_path = None
        video_info = None

        try:
            # 1. 参数校验
            required_params = ['prompt', 'modelType', 'storageService', 'bucketName', 'storedPrefix', 'QueueResult']
            for param in required_params:
                if not task_params.get(param):
                    return {"status": 2, "msg": f"缺少参数: {param}", 'code': 6003}

            model_type = task_params.get('modelType')
            modelName = task_params.get('moduleName')
            # 2. 构建API请求参数与路径
            api_path = ""
            payload = {
                "model": modelName,
                "prompt": task_params.get('prompt'),
                "duration": str(task_params.get('duration', "5")),
                "mode": task_params.get('mode', "std"),
                "cfg_scale": task_params.get('cfg_scale', 0.5),
                "negative_prompt": task_params.get('negative_prompt', '')
            }

            # 根据模型类型补充参数
            if model_type == 1:  # 文生视频
                api_path = "/v1/videos/text2video"
                # payload["camera_control"] = {
                #     "type": "simple",
                #     "config": {"zoom": 5}  # 默认简单运镜
                #     "config": {"zoom": 5}  # 默认简单运镜
                # }
                payload["aspect_ratio"] = task_params.get('ratio', "16:9")

            elif model_type in (2, 4):  # 单图生视频/单图特效
                first_image_url = task_params.get('first_image')
                if not first_image_url:
                    return {"status": 2, "msg": "modelType=2/4需提供first_image", 'code': 6003}
                # 下载并处理首图
                # success, img_data, _ = self._download_and_compress_image(
                #     first_image_url, task_params['storageService'], task_params['bucketName']
                # )
                # if not success:
                #     return {"status": 2, "msg": img_data, 'code': 6005}

                first_image_url = sr_s3.get_url(first_image_url, serviceName=task_params.get('storageService', "s3"), bucketName=task_params.get('bucketName', ""))
                if model_type == 2:
                    api_path = "/v1/videos/image2video"
                    payload["image"] = first_image_url
                else:  # model_type=4
                    api_path = "/v1/videos/effects"
                    payload["effect_scene"] = task_params.get('effect_scene', "default")
                    payload["input"] = task_params.get("inputs", {})
                    payload["image"] = first_image_url

            elif model_type in (3, 5):  # 首尾图生视频/多图特效
                first_image_url = task_params.get('first_image')
                last_image_url = task_params.get('last_image')
                if not (first_image_url and last_image_url):
                    return {"status": 2, "msg": "modelType=3/5需提供first_image和last_image", 'code': 6003}

                first_img = sr_s3.get_url(first_image_url, serviceName=task_params.get('storageService', "s3"), bucketName=task_params.get('bucketName', ""))
                last_img = sr_s3.get_url(last_image_url, serviceName=task_params.get('storageService', "s3"), bucketName=task_params.get('bucketName', ""))
                if model_type == 3:
                    api_path = "/v1/videos/image2video"
                    payload["image"] = first_img
                    payload["image_tail"] = last_img
                else:  # model_type=5
                    api_path = "/v1/videos/effects"
                    payload["effect_scene"] = task_params.get('effect_scene', "default")
                    payload["input"] = task_params.get("inputs", {})
                    payload["images"] = [first_img, last_img]
            else:
                return {"status": 2, "msg": f"不支持的模型类型: {model_type}", 'code': 6003}
            # 3. 创建视频任务

            success, task_api_id, err_msg = self._create_keling_video_task(payload, api_path)
            if not success:
                return {"status": 2, "msg": err_msg, 'code': 6001}
            # task_api_id = '782267034867470403'
            # print(f"任务{task_id}已提交至API，API任务ID: {task_api_id}")

            # 4. 轮询任务结果
            success, video_url = self._poll_task_status(
                task_api_id, api_path,
                timeout=task_params.get('timeout', 600),
                interval=task_params.get('interval', 5)
            )

            if not success:
                if video_url == "Failure to pass the risk control system":
                    return {"status": 7, "msg": video_url, 'code': 7202}
                return {"status": 2, "msg": video_url}
            # 5. 下载视频并上传到S3
            success, video_info, save_path = self._download_video(video_url, task_api_id)
            if not success:
                return {"status": 2, "msg": video_info.get("error"), 'code': 6001}
            # 6. 上传视频到S3
            success, s3_key = self._upload_video_to_s3(
                video_info["save_path"],
                task_params['storedPrefix'],
                task_params['storageService'],
                task_params['bucketName']
            )
            if not success:
                return {"status": 2, "msg": s3_key, 'code': 6001}

            # 7. 发送结果到队列（假设sr_sqs为队列工具）
            from utils import sr_sqs
            msg_body = {
                'status': 1,
                'width': video_info['width'],
                'height': video_info['height'],
                'url': s3_key,
                'local_path': save_path,
                "model_mode": task_params.get('mode', "std"),  # 模型类型标准模式
                'msg': "success",
                'code': ''
            }
            return msg_body

        except Exception as e:
            error_msg = f"任务处理异常: {str(e)}"

            return {"status": 2, "msg": error_msg, 'code': 6001}

        # finally:
        #     # 清理临时文件
        #     self._delete_temp_files(first_image_path,                                    last_image_path,
        #                             video_info.get("save_path") if video_info else None
        #                             )


# --------------------------
# 外部调用示例
# --------------------------
if __name__ == "__main__":
    # 初始化视频生成器
    video_generator = videoGenerator_kling()

    task_params = {
        'taskId': 'e29a9a89b5857d49e8ba',
        'prompt': '跳舞',
        'first_image': 'video/image_to_video/upload/d41d8cd98f00b204e9800998ecf8427e/1754548313665.jfif',
        'last_image': '',
        'duration': '10',
        'fps': '24',
        'resolution': '768p',
        'ratio': '',
        'appId': '1',
        'taskType': 'pwai_video_itf',
        'QueueResult': 'AILocalVideoResult',
        'storageService': 'r2',
        'isLogin': 1,
        'bucketName': 'picwand',
        'storedPrefix': 'video/image_to_video',
        'progress': 1,
        'queueName': 'AIHubVideoTasks_Test',
        'modelType': 2,
        'moduleType': 25,
        'moduleName': 'MiniMax-Hailuo-02',
        'mResolution': '768p',
        'mode': 'std',
        'negative_prompt': '',
        'cfg_scale': 0.5,
        'p_t_b': 0,
        'n_t_b': 0,
    }

    print(task_params)
    # exit()
    # 处理任务
    result = video_generator.process_video_task(task_params)
    print(result)
